"""
Библиотека для отрисовки формул sympy в формате docx.

Взято отсюда:
https://github.com/elapouya/python-docx-template/issues/406#issuecomment-1152999779
"""

from docxtpl import DocxTemplate
from lxml import etree
from sympy import mathml


def _update_get_xml(get_xml_function):
    def get_xml_updated():
        xml = get_xml_function()
        return xml.replace(
            '<m:oMath xmlns:mml="http://www.w3.org/1998/Math/MathML">',
            (
                "<m:oMath xmlns:m="
                '"http://schemas.openxmlformats.org/officeDocument/2006/math">'
            ),
        )

    return get_xml_updated


def _mathml_to_subdoc(mathml_formula: str, docx_template: DocxTemplate):
    tree = etree.fromstring(
        '<math xmlns="http://www.w3.org/1998/Math/MathML">' + mathml_formula + "</math>"
    )
    # convert to MS Office structure
    xslt = etree.parse("utils/xmls/MML2OMML.XSL")
    transform = etree.XSLT(xslt)
    new_dom = transform(tree)
    # write to docx
    equation_document = docx_template.new_subdoc()
    p = equation_document.add_paragraph()
    p._element.append(new_dom.getroot())

    # monkey patch namespace because when adding the subdoc to the template the 'm'
    # namespace gets lost
    equation_document._get_xml = _update_get_xml(equation_document._get_xml)

    return equation_document


def sympy_expr_to_subdoc(expr, docx_template: DocxTemplate):
    return _mathml_to_subdoc(
        mathml(expr, printer="presentation", mul_symbol="dot"),
        docx_template=docx_template,
    )
